package dalle

import (
	"fmt"
	"github.com/go-redis/redis/v8"
	"github.com/imroc/req/v3"
	"gorm.io/gorm"
	"io"
	"myai/core/types"
	logger2 "myai/logger"
	"myai/service"
	"myai/service/oss"
	"myai/store"
	"myai/store/model"
	"myai/utils"
	"time"
)

var log = logger2.GetLogger()

// DALL-E 绘画服务

type Service struct {
	httpClient    *req.Client          // 用于HTTP请求
	db            *gorm.DB             // 数据库操作
	uploadManager *oss.UploaderManager // 文件上传管理器
	taskQueue     *store.RedisQueue    // 任务队列
	notifyQueue   *store.RedisQueue    // 任务通知队列
	userService   *service.UserService
	wsService     *service.WebSocketService // websocket
	clientIds     map[uint]string           // key:jobId value:客户端Id
}

func NewService(db *gorm.DB, manager *oss.UploaderManager, redisCli *redis.Client, userService *service.UserService, wsService *service.WebSocketService) *Service {
	return &Service{
		httpClient:    req.C().SetTimeout(time.Minute * 3),
		db:            db,
		taskQueue:     store.NewRedisQueue("My_DallE_Task_Queue", redisCli),
		notifyQueue:   store.NewRedisQueue("My_DallE_Notify_Queue", redisCli),
		wsService:     wsService,
		uploadManager: manager,
		userService:   userService,
		clientIds:     map[uint]string{},
	}
}

// PushTask 向Redis队列中推送任务
func (s *Service) PushTask(task types.DallTask) {
	log.Infof("将新的DALL-E任务添加到任务列表中：%+v", task)
	s.taskQueue.RPush(task)
}

func (s *Service) Run() {
	// 将数据库未提交的任务加载到队列
	var jobs []model.DallJob
	s.db.Where("progress", 0).Find(&jobs)
	for _, v := range jobs {
		var task types.DallTask
		err := utils.JsonDecode(v.TaskInfo, &task)
		if err != nil {
			log.Errorf("解码任务信息时出错：%v", err)
			continue
		}
		task.Id = v.Id
		// 推送任务到队列中
		s.PushTask(task)
	}
	log.Info("启动DALL-E任务消费者 ...")
	go func() {
		for {
			var task types.DallTask
			// 消费任务
			err := s.taskQueue.LPop(&task)
			if err != nil {
				log.Errorf("消费任务时出错：%v", err)
				continue
			}
			log.Infof("处理新的 DALL-E 任务：%+v", task)
			// 存储任务和客户端的对应关系
			s.clientIds[task.Id] = task.ClientId
			// 开始执行绘图任务
			_, err = s.Image(task, false)
			if err != nil {
				log.Errorf("绘图任务执行失败：%v", err)
				// 修改任务状态
				s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(map[string]interface{}{
					"progress": service.FailTaskProgress,
					"err_msg":  err.Error(),
				})
				// 发送通知消息
				s.notifyQueue.RPush(service.NotifyMessage{
					ClientId: task.ClientId,
					UserId:   int(task.UserId),
					JobId:    int(task.Id),
					Message:  service.TaskStatusFailed,
				})
			}
		}
	}()
}

// imgReq 请求 DALLE 模型参数
type imgReq struct {
	Model   string `json:"model"`             // 模型
	Prompt  string `json:"prompt"`            // 提示词
	N       int    `json:"n,omitempty"`       // 生成图像数量 生成图像数量。DALL-E 3仅支持单次生成1张图像（DALL-E 2支持最多10张）。
	Size    string `json:"size,omitempty"`    // 图像尺寸 1024x1024（默认） 1792x1024（宽屏） 1024x1792（高屏）
	Quality string `json:"quality,omitempty"` // 图像质量 standard（标准） hd（高清）
	Style   string `json:"style,omitempty"`   // 风格选择：vivid强调戏剧化效果（超现实风格），natural更接近真实摄影效果
}

// imgRes DALLE 模型响应结果
type imgRes struct {
	Created int64 `json:"created"` // 图像生成时间（UTC）。
	Data    []struct {
		RevisedPrompt string `json:"revised_prompt,omitempty"` // revised_prompt：优化后的提示语（系统自动调整）
		Url           string `json:"url,omitempty"`            // 临时访问链接
		B64Json       string `json:"b64_json,omitempty"`
	} `json:"data"`
}

// ErrRes 错误响应结果
type ErrRes struct {
	Error struct {
		Code    interface{} `json:"code"`
		Message string      `json:"message"`
		Param   interface{} `json:"param"`
		Type    string      `json:"type"`
	} `json:"error"`
}

// Image dall-e进行绘图
func (s *Service) Image(task types.DallTask, sync bool) (string, error) {
	// 1、把中文提示词翻译为英文
	// 2、查询用户信息，判断算力是否足够
	// 3、扣减算力
	// 4、查询DALL-E 的 API KEY
	// 5、发送绘图请求
	// 6、更新API KEY最后一次使用的时间
	// 7、请求完成 更新 dall 任务的进度
	// 8、向通知队列中发送消息，返回给客户端用户
	// 9、下载图片返回
	log.Debugf("绘图参数：%+v", task)
	prompt := task.Prompt
	// 判断提示词中是否包含中文，如果有中文就翻译提示词
	if utils.HasChinese(prompt) {
		// 提示词翻译成英文
		content, err := utils.OpenAIRequest(s.db, fmt.Sprintf(service.TranslatePromptTemplate, prompt), task.TranslateModelId)
		if err == nil {
			prompt = content
			log.Debugf("重写后的提示词：%s", prompt)
		}
	}

	// 查询模型信息
	var chatModel model.ChatModel
	s.db.Where("id = ?", task.ModelId).First(&chatModel)

	// 查询 DALLE API-KEY
	var apiKey model.ApiKey
	session := s.db.Where("enabled", true)
	if chatModel.KeyId > 0 {
		session = session.Where("id = ?", chatModel.KeyId)
	} else {
		session = session.Where("type = ?", "dalle")
	}
	err := session.Order("last_used_at asc").First(&apiKey).Error
	if err != nil {
		return "", fmt.Errorf("没有可用的绘图模型的密钥：%v", err)
	}

	var res imgRes
	var errRes ErrRes
	if len(apiKey.ProxyURL) > 5 {
		s.httpClient.SetProxyURL(apiKey.ProxyURL).R()
	}
	apiURL := fmt.Sprintf("%s/v1/images/generations", apiKey.ApiURL)
	reqBody := imgReq{
		Model:   chatModel.Value,
		Prompt:  prompt,
		N:       1,
		Size:    task.Size,
		Style:   task.Style,
		Quality: task.Quality,
	}
	log.Infof("Dall 请求地址：%s，API KEY：%s，请求参数：%+v", apiURL, apiKey.Value, reqBody)
	// 发送请求
	r, err := s.httpClient.R().SetHeader("Body-Type", "application/json").
		SetHeader("Authorization", "Bearer "+apiKey.Value).
		SetBody(reqBody).
		SetErrorResult(&errRes).
		SetSuccessResult(&res).
		Post(apiURL)

	if err != nil {
		return "", fmt.Errorf("发送请求时出错：%v", err)
	}
	// 判断HTTP状态码是否>=400，如果是就返回true
	if r.IsErrorState() {
		return "", fmt.Errorf("发送请求，状态错误: %s, %+v", r.Status, errRes.Error)
	}
	// 读取结果
	all, _ := io.ReadAll(r.Body)
	log.Debugf("响应结果：%+v", string(all))

	// 更新API KEY最后一次使用的时间
	s.db.Model(&apiKey).UpdateColumn("last_used_at", time.Now().Unix())

	var imgURL string
	var data = map[string]interface{}{
		"progress": 100,
		"prompt":   prompt,
	}

	// 如果返回的事base64，则需要上传到OSS
	if res.Data[0].B64Json != "" {
		imgURL, err = s.uploadManager.GetUploadHandler().PutBase64(res.Data[0].B64Json)
		if err != nil {
			return "", fmt.Errorf("把生成的Base64图片上传到OSS失败：%v", err)
		}
		log.Infof("上传图片到OSS：%s", imgURL)
		data["img_url"] = imgURL
	} else {
		imgURL = res.Data[0].Url
	}
	data["org_url"] = imgURL

	// 更新 dall 任务的进度
	err = s.db.Model(&model.DallJob{Id: task.Id}).UpdateColumns(data).Error

	if err != nil {
		return "", fmt.Errorf("dalle 更新数据库任务时出错：%v", err)
	}

	// 向通知队列中发送消息
	s.notifyQueue.RPush(
		service.NotifyMessage{
			ClientId: task.ClientId,
			UserId:   int(task.UserId),
			JobId:    int(task.Id),
			Message:  service.TaskStatusFailed,
		})
	var content string
	if sync {
		imgURL, err := s.downloadImage(task.Id, int(task.UserId), res.Data[0].Url)
		if err != nil {
			return "", fmt.Errorf("下载图片出错：%v", err)
		}
		// markdown 语法返回
		content = fmt.Sprintf("```\n%s\n```\n下面是我为你创作的图片：\n\n![](%s)\n", prompt, imgURL)
	}
	return content, nil
}

// CheckTaskNotify 对通知队列内的数据进行处理，发送给客户端用户
func (s *Service) CheckTaskNotify() {
	go func() {
		log.Infof("正在运行DALL-E任务通知检查 ...")
		for {
			var message service.NotifyMessage
			// 从队列中消费数据
			err := s.notifyQueue.LPop(&message)
			if err != nil {
				log.Errorf("notifyQueue 队列消费数据失败：%v", err)
				continue
			}
			log.Debugf("通知消息：%+v", message)
			// 获取websocket连接
			client := s.wsService.Clients.Get(message.ClientId)
			if client == nil {
				log.Warn("获取到的客户端连接为空")
				continue
			}
			// 向客户端发送消息
			utils.SendChannelMsg(client, types.ChDall, message.Message)
		}
	}()
}

// CheckTaskStatus 检查任务状态
func (s *Service) CheckTaskStatus() {
	go func() {
		log.Info("正在运行DALL-E任务状态检查 ...")
		for {
			// 检查未完成任务进度
			var jobs []model.DallJob
			s.db.Where("progress < ?", 100).Find(&jobs)
			for _, job := range jobs {
				// 超时任务标记失败
				// 检查当前时间是否超过任务创建时间10分钟以上
				// Sub方法应该是用来计算两个时间点之间的差值。这个方法会返回一个Duration类型的值，表示两个时间之间的间隔。
				if time.Now().Sub(job.CreatedAt) > time.Minute*10 {
					job.Progress = service.FailTaskProgress
					job.ErrMsg = "任务超时"
					// 更新任务
					s.db.Updates(&job)
				}
			}
			// 找出失败任务，并恢复其扣减的算力
			s.db.Where("progress", service.FailTaskProgress).Where("power > ?", 0).Find(&jobs)
			for _, job := range jobs {
				var task types.DallTask
				err := utils.JsonDecode(job.TaskInfo, &task)
				if err != nil {
					log.Errorf("dalle CheckTaskStatus JsonDecode 出现错误：%+v", err)
					continue
				}
				err = s.userService.IncreasePower(int(job.UserId), job.Power, model.PowerLog{
					Type:   types.PowerRefund,
					Model:  task.ModelName,
					Remark: fmt.Sprintf("任务失败，退回算力。任务Id：%d，Err：%s", job.Id, job.ErrMsg),
				})
				if err != nil {
					log.Errorf("dall 任务执行失败，退回算力时操作数据库失败：%v", err)
					continue
				}
				// 更新任务状态
				s.db.Model(&job).UpdateColumn("power", 0)
			}
			time.Sleep(time.Second * 10)
		}
	}()
}

// DownloadImages 批量下载所有完成任务的图片
func (s *Service) DownloadImages() {
	go func() {
		var items []model.DallJob
		for {
			// 查询任务进度为百分之百且没有图片的任务
			res := s.db.Where("img_url = ? and progress = ?", "", 100).Find(&items)
			if res.Error != nil {
				log.Errorf("批量下载完成任务的图片时，查询数据库出错：%v", res.Error.Error())
				continue
			}
			for _, v := range items {
				if v.OrgURL == "" {
					continue
				}
				log.Infof("尝试下载图像：%s", v.OrgURL)
				// 下载图片
				imgURL, err := s.downloadImage(v.Id, int(v.UserId), v.OrgURL)
				if err != nil {
					log.Errorf("下载图片时出错：%s，错误：%v", imgURL, err)
					continue
				} else {
					log.Infof("成功下载图片：%s", imgURL)
				}
			}
			time.Sleep(time.Second * 6)
		}
	}()
}

// downloadImage 下载图片
func (s *Service) downloadImage(jobId uint, userId int, orgURL string) (string, error) {
	// 下载图片到指定位置
	imgURL, err := s.uploadManager.GetUploadHandler().PutUrlFile(orgURL, false)
	if err != nil {
		return "", err
	}
	// 更新图片地址
	res := s.db.Model(&model.DallJob{Id: jobId}).UpdateColumn("img_url", imgURL)
	if res.Error != nil {
		return "", err
	}
	// 发送通知
	s.notifyQueue.RPush(service.NotifyMessage{ClientId: s.clientIds[jobId], UserId: userId, JobId: int(jobId), Message: service.TaskStatusFinished})
	return imgURL, nil
}
